import os

import pandas as pd
import numpy as np
from sklearn.preprocessing import MinMaxScaler
from datetime import date, timedelta

from model.prediction import transformer_model
from utils.data_source import query, engine, update
from utils.data_utils import query_hist

# 归一化是方便模型计算，反归一化可以把值还原
MMS = MinMaxScaler()


# 创意一可以被训练的序列
def create_sequence(dataset, window_size=11):
    sequences = list()
    labels = list()
    dataset_th = len(dataset)
    for idx in range(dataset_th):
        if idx + window_size >= dataset_th:
            break
        sequences.append(dataset.iloc[idx:idx + window_size])  # 输入窗口
        labels.append(dataset.iloc[idx + window_size])  # 输出值（下一时间步）
    return np.array(sequences), np.array(labels)


def prediction(dataa):
    # 查詢股票歷史
    hist, label = query_hist(dataa=dataa)
    # Extracting required columns开盘,收盘,最高,最低
    hist.set_index('data', drop=True, inplace=True)
    hist.sort_index(inplace=True)
    hist[hist.columns] = MMS.fit_transform(hist)
    # 取前80%的数据作为训练样本，后20作为测试样本
    training_size = round(len(hist) * 0.80)
    train_data, test_data = hist[:training_size], hist[training_size:]
    train_seq, train_label = create_sequence(train_data)
    test_seq, test_label = create_sequence(test_data)
    values = {}
    prediction_data = list()
    run_num = 10
    for idx in range(run_num):
        upcoming_prediction = pd.DataFrame(
            columns=[label],
            index=pd.date_range(start=hist.index[-1] + timedelta(days=1), periods=1))
        upcoming_prediction.index = pd.to_datetime(upcoming_prediction.index)
        transformer_model(train_seq, train_label, test_seq, test_label, upcoming_prediction)
        prediction_item = MMS.inverse_transform(upcoming_prediction[label])
        for index in range(len(label)):
            col_name = label[index]
            if 'close' in col_name:
                symbol = col_name.strip("close")
                if symbol in values.keys():
                    values[symbol] = prediction_item[0][index] + values[symbol]
                else:
                    values[symbol] = prediction_item[0][index]
    hist[hist.columns] = MMS.inverse_transform(hist)
    time = hist.index[-1] + timedelta(days=1)
    time_str = time.__format__("%m%d")
    for symbol in values.keys():
        val = hist.iloc[-1][symbol + 'close']
        value = (values[symbol] / run_num - val) * 100 / val
        prediction_data.append([time, symbol, value])
    csv_date = pd.DataFrame(data=prediction_data, columns=["时间", "涨跌排序", "涨跌排序zdf"]).sort_values(
        by="涨跌排序zdf", ascending=False)
    try:
        update(f"""delete from world.xl_result where 时间 = '{time.__format__("%Y%m%d")}'""")
    except Exception as r:
        print(r)
    csv_date.to_sql(f'xl_result', con=engine(), if_exists='append', index=False)
    return csv_date.to_csv(f"C:\\Users\\pc\\Desktop\\结果{time_str}.csv")
